import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

# 数据预处理
def load_data():
    base = pd.read_csv('customer_base.csv')
    behavior = pd.read_csv('customer_behavior_assets.csv')
    
    # 获取最新的客户资产数据（按统计月份排序，取每个客户的最新记录）
    latest_data = behavior.sort_values('stat_month').groupby('customer_id').tail(1)
    
    # 合并基础信息和行为资产数据
    merged = pd.merge(base, latest_data, on='customer_id', how='inner')
    
    # 计算客户未来3个月资产达到100万+的概率
    # 假设资产线性增长，基于最近两个月的数据计算增长率
    customer_sorted = behavior.sort_values(['customer_id', 'stat_month'])
    customer_sorted['prev_assets'] = customer_sorted.groupby('customer_id')['total_assets'].shift(1)
    customer_sorted['asset_growth'] = customer_sorted['total_assets'] - customer_sorted['prev_assets']
    customer_sorted['growth_rate'] = customer_sorted['asset_growth'] / customer_sorted['prev_assets']

    # 计算每个客户的平均增长率
    growth_rates = customer_sorted.groupby('customer_id')['growth_rate'].mean().reset_index()
    growth_rates.columns = ['customer_id', 'avg_growth_rate']

    # 合并增长率数据
    merged = pd.merge(merged, growth_rates, on='customer_id', how='left')

    # 填充缺失的增长率（对于只有一条记录的客户）
    merged['avg_growth_rate'] = merged['avg_growth_rate'].fillna(0)

    # 计算3个月后预测资产
    merged['predicted_assets_3m'] = merged['total_assets'] * (1 + merged['avg_growth_rate']) ** 3

    # 创建目标变量：3个月内资产是否能达到100万+
    merged['target'] = (merged['predicted_assets_3m'] >= 1000000).astype(int)
    
    # 选择关键特征
    features = merged[[
        'age', 'monthly_income', 'total_assets', 'deposit_balance', 'financial_balance', 
        'fund_balance', 'insurance_balance', 'product_count', 'financial_repurchase_count',
        'credit_card_monthly_expense', 'investment_monthly_count', 'app_login_count',
        'app_financial_view_time', 'app_product_compare_count', 'avg_growth_rate'
    ]]
    
    return features, merged['target']

# 模型训练与可视化
def train_and_visualize():
    X, y = load_data()
    
    # 处理缺失值
    X = X.fillna(0)
    
    # 标准化连续特征
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)
    
    # 训练逻辑回归
    lr = LogisticRegression(max_iter=2000)
    lr.fit(X_train, y_train)
    
    # 生成系数可视化
    plt.figure(figsize=(10, 8))
    coef_df = pd.DataFrame({
        'feature': X.columns,
        'coefficient': lr.coef_[0]
    }).sort_values('coefficient', key=abs, ascending=False)
    
    colors = ['blue' if x > 0 else 'red' for x in coef_df['coefficient']]
    plt.barh(range(len(coef_df)), coef_df['coefficient'], color=colors)
    plt.yticks(range(len(coef_df)), coef_df['feature'])
    plt.xlabel('Coefficient Value')
    plt.title('Logistic Regression Coefficients for Predicting 1M+ Assets in 3 Months')
    plt.axvline(x=0, color='black', linewidth=0.5)
    
    # 添加图例
    red_patch = plt.Rectangle((0,0),1,1, color='red')
    blue_patch = plt.Rectangle((0,0),1,1, color='blue')
    plt.legend([blue_patch, red_patch], ['Positive Coefficient', 'Negative Coefficient'])
    
    plt.tight_layout()
    plt.savefig('image_show/lr_coefficients.png', dpi=300, bbox_inches='tight')
    plt.close()
    
    print("逻辑回归系数：")
    print(coef_df)
    print(f"\n模型准确率: {lr.score(X_test, y_test):.4f}")
    
    # 保存系数到CSV文件
    coef_df.to_csv('image_show/lr_coefficients.csv', index=False)
    print("\n逻辑回归系数已保存到 lr_coefficients.csv 文件")

if __name__ == '__main__':
    train_and_visualize()